{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/google/neural-tangents/blob/main/notebooks/experimental/empirical_ntk_resnet_tf.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nTt0UNQbk_Td"
   },
   "source": [
    "# Example of computing NTK of a **Tensorflow (Keras)** ResNet50 on ImageNet inputs\n",
    "Warning: computing the NTK in Tensorflow currently appears to have very long compile times (but OK runtime), can be prone to triggering XLA errors, and does not distinguish between trainable and non-trainable parameters of the model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dvXMUSdFjCqq"
   },
   "source": [
    "Tested on NVIDIA A100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oxdetCnZH7xE"
   },
   "source": [
    "More examples: \n",
    "\n",
    "\n",
    "*   JAX (Flax):\n",
    "  * [FCN](https://colab.research.google.com/github/google/neural-tangents/blob/main/notebooks/empirical_ntk_fcn.ipynb)\n",
    "  * [ResNet18](https://colab.research.google.com/github/google/neural-tangents/blob/main/notebooks/empirical_ntk_resnet.ipynb)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Sl2JyRE1hK-z"
   },
   "source": [
    "# Imports and setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Xx1xr9v5pyl0",
    "outputId": "754443ca-7836-483f-f40f-7513e6e61c04"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "GPU 0: A100-SXM4-40GB (UUID: GPU-07b846cf-7f39-fff7-8224-d367cef00104)\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi -L"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ZlckZChsxTWj",
    "outputId": "5c6d6cdb-c398-46bc-fd8e-b73c0807d0fe"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Requirement already satisfied: pip in /usr/local/lib/python3.7/dist-packages (21.1.3)\n",
      "Collecting pip\n",
      "  Downloading pip-22.1.2-py3-none-any.whl (2.1 MB)\n",
      "\u001B[K     |████████████████████████████████| 2.1 MB 6.7 MB/s \n",
      "\u001B[?25hInstalling collected packages: pip\n",
      "  Attempting uninstall: pip\n",
      "    Found existing installation: pip 21.1.3\n",
      "    Uninstalling pip-21.1.3:\n",
      "      Successfully uninstalled pip-21.1.3\n",
      "Successfully installed pip-22.1.2\n",
      "Looking in links: https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
      "Requirement already satisfied: jax[cuda11_cudnn805] in /usr/local/lib/python3.7/dist-packages (0.3.1)\n",
      "Collecting jax[cuda11_cudnn805]\n",
      "  Downloading jax-0.3.13.tar.gz (951 kB)\n",
      "\u001B[2K     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m951.0/951.0 kB\u001B[0m \u001B[31m26.4 MB/s\u001B[0m eta \u001B[36m0:00:00\u001B[0m\n",
      "\u001B[?25h  Preparing metadata (setup.py) ... \u001B[?25l\u001B[?25hdone\n",
      "Requirement already satisfied: absl-py in /usr/local/lib/python3.7/dist-packages (from jax[cuda11_cudnn805]) (1.0.0)\n",
      "Requirement already satisfied: numpy>=1.19 in /usr/local/lib/python3.7/dist-packages (from jax[cuda11_cudnn805]) (1.21.5)\n",
      "Requirement already satisfied: opt_einsum in /usr/local/lib/python3.7/dist-packages (from jax[cuda11_cudnn805]) (3.3.0)\n",
      "Requirement already satisfied: scipy>=1.2.1 in /usr/local/lib/python3.7/dist-packages (from jax[cuda11_cudnn805]) (1.4.1)\n",
      "Requirement already satisfied: typing_extensions in /usr/local/lib/python3.7/dist-packages (from jax[cuda11_cudnn805]) (3.10.0.2)\n",
      "Collecting jaxlib==0.3.10+cuda11.cudnn805\n",
      "  Downloading https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.3.10%2Bcuda11.cudnn805-cp37-none-manylinux2014_x86_64.whl (175.7 MB)\n",
      "\u001B[2K     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m175.7/175.7 MB\u001B[0m \u001B[31m8.4 MB/s\u001B[0m eta \u001B[36m0:00:00\u001B[0m\n",
      "\u001B[?25hRequirement already satisfied: flatbuffers<3.0,>=1.12 in /usr/local/lib/python3.7/dist-packages (from jaxlib==0.3.10+cuda11.cudnn805->jax[cuda11_cudnn805]) (2.0)\n",
      "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from absl-py->jax[cuda11_cudnn805]) (1.15.0)\n",
      "Building wheels for collected packages: jax\n",
      "  Building wheel for jax (setup.py) ... \u001B[?25l\u001B[?25hdone\n",
      "  Created wheel for jax: filename=jax-0.3.13-py3-none-any.whl size=1099581 sha256=102337595723f392fa5580580f6ccea4c38091fa02d9c9e2d0346d5fb0fabad2\n",
      "  Stored in directory: /root/.cache/pip/wheels/c4/6e/2c/605117cd74ef8db71f2c9e4b059b73f03c091546388d015d9f\n",
      "Successfully built jax\n",
      "Installing collected packages: jaxlib, jax\n",
      "  Attempting uninstall: jaxlib\n",
      "    Found existing installation: jaxlib 0.3.0+cuda11.cudnn805\n",
      "    Uninstalling jaxlib-0.3.0+cuda11.cudnn805:\n",
      "      Successfully uninstalled jaxlib-0.3.0+cuda11.cudnn805\n",
      "  Attempting uninstall: jax\n",
      "    Found existing installation: jax 0.3.1\n",
      "    Uninstalling jax-0.3.1:\n",
      "      Successfully uninstalled jax-0.3.1\n",
      "Successfully installed jax-0.3.13 jaxlib-0.3.10+cuda11.cudnn805\n",
      "\u001B[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001B[0m\u001B[33m\n",
      "\u001B[0mCollecting git+https://www.github.com/google/neural-tangents.git\n",
      "  Cloning https://www.github.com/google/neural-tangents.git to /tmp/pip-req-build-nep6el53\n",
      "  Running command git clone --filter=blob:none --quiet https://www.github.com/google/neural-tangents.git /tmp/pip-req-build-nep6el53\n",
      "  warning: redirecting to https://github.com/google/neural-tangents.git/\n",
      "  Resolved https://www.github.com/google/neural-tangents.git to commit eec18f56562864540f8141372d7df67f9c34373e\n",
      "  Preparing metadata (setup.py) ... \u001B[?25l\u001B[?25hdone\n",
      "Requirement already satisfied: jax>=0.3.13 in /usr/local/lib/python3.7/dist-packages (from neural-tangents==0.6.0) (0.3.13)\n",
      "Collecting frozendict>=2.3\n",
      "  Downloading frozendict-2.3.2-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (99 kB)\n",
      "\u001B[2K     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m99.0/99.0 kB\u001B[0m \u001B[31m6.1 MB/s\u001B[0m eta \u001B[36m0:00:00\u001B[0m\n",
      "\u001B[?25hCollecting typing_extensions>=4.0.1\n",
      "  Downloading typing_extensions-4.2.0-py3-none-any.whl (24 kB)\n",
      "Collecting tf2jax>=0.3.0\n",
      "  Downloading tf2jax-0.3.0-py3-none-any.whl (63 kB)\n",
      "\u001B[2K     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m63.1/63.1 kB\u001B[0m \u001B[31m11.4 MB/s\u001B[0m eta \u001B[36m0:00:00\u001B[0m\n",
      "\u001B[?25hRequirement already satisfied: scipy>=1.2.1 in /usr/local/lib/python3.7/dist-packages (from jax>=0.3.13->neural-tangents==0.6.0) (1.4.1)\n",
      "Requirement already satisfied: opt-einsum in /usr/local/lib/python3.7/dist-packages (from jax>=0.3.13->neural-tangents==0.6.0) (3.3.0)\n",
      "Requirement already satisfied: absl-py in /usr/local/lib/python3.7/dist-packages (from jax>=0.3.13->neural-tangents==0.6.0) (1.0.0)\n",
      "Requirement already satisfied: numpy>=1.19 in /usr/local/lib/python3.7/dist-packages (from jax>=0.3.13->neural-tangents==0.6.0) (1.21.5)\n",
      "Requirement already satisfied: jaxlib>=0.1.37 in /usr/local/lib/python3.7/dist-packages (from tf2jax>=0.3.0->neural-tangents==0.6.0) (0.3.10+cuda11.cudnn805)\n",
      "Requirement already satisfied: tensorflow>=2.8.0 in /usr/local/lib/python3.7/dist-packages (from tf2jax>=0.3.0->neural-tangents==0.6.0) (2.8.0)\n",
      "Requirement already satisfied: dm-tree>=0.1.5 in /usr/local/lib/python3.7/dist-packages (from tf2jax>=0.3.0->neural-tangents==0.6.0) (0.1.6)\n",
      "Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.7/dist-packages (from dm-tree>=0.1.5->tf2jax>=0.3.0->neural-tangents==0.6.0) (1.15.0)\n",
      "Requirement already satisfied: flatbuffers<3.0,>=1.12 in /usr/local/lib/python3.7/dist-packages (from jaxlib>=0.1.37->tf2jax>=0.3.0->neural-tangents==0.6.0) (2.0)\n",
      "Requirement already satisfied: libclang>=9.0.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (13.0.0)\n",
      "Requirement already satisfied: keras-preprocessing>=1.1.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (1.1.2)\n",
      "Collecting tf-estimator-nightly==2.8.0.dev2021122109\n",
      "  Downloading tf_estimator_nightly-2.8.0.dev2021122109-py2.py3-none-any.whl (462 kB)\n",
      "\u001B[2K     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m462.5/462.5 kB\u001B[0m \u001B[31m25.1 MB/s\u001B[0m eta \u001B[36m0:00:00\u001B[0m\n",
      "\u001B[?25hRequirement already satisfied: gast>=0.2.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (0.5.3)\n",
      "Requirement already satisfied: wrapt>=1.11.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (1.13.3)\n",
      "Requirement already satisfied: protobuf>=3.9.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (3.17.3)\n",
      "Requirement already satisfied: keras<2.9,>=2.8.0rc0 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (2.8.0)\n",
      "Requirement already satisfied: astunparse>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (1.6.3)\n",
      "Requirement already satisfied: google-pasta>=0.1.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (0.2.0)\n",
      "Requirement already satisfied: tensorboard<2.9,>=2.8 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (2.8.0)\n",
      "Requirement already satisfied: tensorflow-io-gcs-filesystem>=0.23.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (0.24.0)\n",
      "Requirement already satisfied: setuptools in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (57.4.0)\n",
      "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (1.1.0)\n",
      "Requirement already satisfied: h5py>=2.9.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (3.1.0)\n",
      "Requirement already satisfied: grpcio<2.0,>=1.24.3 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (1.44.0)\n",
      "Requirement already satisfied: wheel<1.0,>=0.23.0 in /usr/local/lib/python3.7/dist-packages (from astunparse>=1.6.0->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (0.37.1)\n",
      "Requirement already satisfied: cached-property in /usr/local/lib/python3.7/dist-packages (from h5py>=2.9.0->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (1.5.2)\n",
      "Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (0.6.1)\n",
      "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (1.8.1)\n",
      "Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (1.0.1)\n",
      "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (0.4.6)\n",
      "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (3.3.6)\n",
      "Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (2.23.0)\n",
      "Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (1.35.0)\n",
      "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.7/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (4.8)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.7/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (0.2.8)\n",
      "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (4.2.4)\n",
      "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.7/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (1.3.1)\n",
      "Requirement already satisfied: importlib-metadata>=4.4 in /usr/local/lib/python3.7/dist-packages (from markdown>=2.6.8->tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (4.11.1)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (2021.10.8)\n",
      "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (2.10)\n",
      "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (3.0.4)\n",
      "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (1.24.3)\n",
      "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata>=4.4->markdown>=2.6.8->tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (3.7.0)\n",
      "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.7/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (0.4.8)\n",
      "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.7/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.9,>=2.8->tensorflow>=2.8.0->tf2jax>=0.3.0->neural-tangents==0.6.0) (3.2.0)\n",
      "Building wheels for collected packages: neural-tangents\n",
      "  Building wheel for neural-tangents (setup.py) ... \u001B[?25l\u001B[?25hdone\n",
      "  Created wheel for neural-tangents: filename=neural_tangents-0.6.0-py3-none-any.whl size=242807 sha256=717150e6d46b4f7e2e39fc6fd5355cb9043ac81e69eec983c6491eb3fac14484\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-7wgkgyd1/wheels/c4/d7/cf/1d7c8b0f6181840ab3f5ddf56850700097cd4582c149fa2877\n",
      "Successfully built neural-tangents\n",
      "Installing collected packages: tf-estimator-nightly, typing_extensions, frozendict, tf2jax, neural-tangents\n",
      "  Attempting uninstall: typing_extensions\n",
      "    Found existing installation: typing-extensions 3.10.0.2\n",
      "    Uninstalling typing-extensions-3.10.0.2:\n",
      "      Successfully uninstalled typing-extensions-3.10.0.2\n",
      "\u001B[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "arviz 0.11.4 requires typing-extensions<4,>=3.7.4.3, but you have typing-extensions 4.2.0 which is incompatible.\u001B[0m\u001B[31m\n",
      "\u001B[0mSuccessfully installed frozendict-2.3.2 neural-tangents-0.6.0 tf-estimator-nightly-2.8.0.dev2021122109 tf2jax-0.3.0 typing_extensions-4.2.0\n",
      "\u001B[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001B[0m\u001B[33m\n",
      "\u001B[0m"
     ]
    }
   ],
   "source": [
    "!pip install --upgrade pip\n",
    "!pip install --upgrade jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
    "!pip install git+https://www.github.com/google/neural-tangents.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "LbW8KVnsPfVd"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import neural_tangents as nt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "CqxnhMKDE2Gf"
   },
   "outputs": [],
   "source": [
    "input_shape = (224, 224, 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eqVgigfhDoHU"
   },
   "source": [
    "# Tensorflow model definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "wdqBa3OQDf97"
   },
   "outputs": [],
   "source": [
    "def get_model(O: int) -> tf.Module:\n",
    "  return tf.keras.applications.resnet.ResNet50(classes=O, weights=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oXqlToEouqSc"
   },
   "source": [
    "# NTK functions declaration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "lPh5LGz9JBK_"
   },
   "outputs": [],
   "source": [
    "def get_ntk_fns(O: int):\n",
    "  # Define a TF-Keras ResNet50 with `O` output logits.\n",
    "  f = get_model(O)\n",
    "  f.build((None, *input_shape))\n",
    "  _, params = nt.experimental.get_apply_fn_and_params(f)\n",
    "\n",
    "  kwargs = dict(\n",
    "      f=f,\n",
    "      trace_axes=(),\n",
    "      vmap_axes=0\n",
    "  )\n",
    "\n",
    "  # Different NTK implementations\n",
    "  jacobian_contraction = nt.experimental.empirical_ntk_fn_tf(\n",
    "      **kwargs, implementation=nt.NtkImplementation.JACOBIAN_CONTRACTION)\n",
    "  ntvp = nt.experimental.empirical_ntk_fn_tf(\n",
    "      **kwargs, implementation=nt.NtkImplementation.NTK_VECTOR_PRODUCTS)\n",
    "  str_derivatives = nt.experimental.empirical_ntk_fn_tf(\n",
    "      **kwargs, implementation=nt.NtkImplementation.STRUCTURED_DERIVATIVES)\n",
    "  auto = nt.experimental.empirical_ntk_fn_tf(\n",
    "      **kwargs, implementation=nt.NtkImplementation.AUTO)\n",
    "  \n",
    "  return params, (jacobian_contraction, ntvp, str_derivatives, auto)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0lWFC3QEgao4"
   },
   "source": [
    "# $\\color{blue}O = 8$ logit, batch size $\\color{red}N = 8$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kbExWKkUg9ew"
   },
   "source": [
    "Structured derivatives compute NTK fastest. NTK-vector products are actually slower in this setting, due to costly forward pass relative to parameters size, and therefore scales poorly with batch size $\\color{red}N$. While it scales better with $\\color{blue}O$ than other methods, it's not enough to overcome the $\\color{red}N^2$ forward passes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "bwhIZWqxKTlt",
    "outputId": "ca35c5c4-b0d1-4e13-9bc1-4e033e1ef472"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "/usr/local/lib/python3.7/dist-packages/neural_tangents/experimental/empirical_tf/empirical.py:232: UserWarning: This function is an early proof-of-concept.\n",
      "  warnings.warn('This function is an early proof-of-concept.')\n"
     ]
    }
   ],
   "source": [
    "O = 8\n",
    "N = 8\n",
    "\n",
    "# Input images x\n",
    "x1 = tf.random.normal((N, *input_shape))\n",
    "x2 = tf.random.normal((N, *input_shape))\n",
    "\n",
    "params, (ntk_fn_jacobian_contraction, ntk_fn_ntvp, ntk_fn_str_derivatives, ntk_fn_auto) = get_ntk_fns(O=O)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "zObT8WnPggFo",
    "outputId": "e91135d3-49bc-42c6-c92c-24cf57d777f3"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(8, 8, 8, 8)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Jacobian contraction\n",
    "k_1 = ntk_fn_jacobian_contraction(x1, x2, params)\n",
    "print(k_1.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "FW9gJJ4qggFp",
    "outputId": "57341f6f-6bbd-4e3d-82bb-9c92b85f3764"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(8, 8, 8, 8)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# NTK-vector products\n",
    "k_2 = ntk_fn_ntvp(x1, x2, params)\n",
    "print(k_2.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "gFeWnqGQggFp",
    "outputId": "3dd588cf-fefa-49d4-eb48-553965f9174e"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(8, 8, 8, 8)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Structured derivatives\n",
    "k_3 = ntk_fn_str_derivatives(x1, x2, params)\n",
    "print(k_3.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Q63v1L1aggFp",
    "outputId": "f6458455-6cb7-496b-f6ca-c67988bb56ba"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "tf.Tensor(0.00030575716, shape=(), dtype=float32) tf.Tensor(0.00075477577, shape=(), dtype=float32) tf.Tensor(0.0010306665, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Make sure kernels agree.\n",
    "print(\n",
    "    tf.reduce_max(tf.abs(k_1 - k_2)) / tf.reduce_mean(tf.abs(k_1)), \n",
    "    tf.reduce_max(tf.abs(k_1 - k_3)) / tf.reduce_mean(tf.abs(k_1)),\n",
    "    tf.reduce_max(tf.abs(k_2 - k_3)) / tf.reduce_mean(tf.abs(k_2))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "gMsFPSOr3DTE",
    "outputId": "259a1982-2eca-461f-ac12-882db0100e74"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "impl=1, flops=17815195648.0\n",
      "impl=2, flops=61376139264.0\n",
      "impl=3, flops=17957609472.0\n",
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(8, 8, 8, 8)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Selects best method based on FLOPs at first call / compilation.\n",
    "# Takes about 3x more time to compile.\n",
    "# WARNING: due to an XLA issue, currently only works correctly on TPUs!\n",
    "# Wrong FLOPs for CPU/GPU of JITted functions.\n",
    "k_0 = ntk_fn_auto(x1, x2, params)\n",
    "print(k_0.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "diP7nkBuggFp",
    "outputId": "1ab7f2a2-7da1-4559-bdf4-9a173a04d380"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "1 loop, best of 5: 319 ms per loop\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "%%timeit\n",
    "ntk_fn_jacobian_contraction(x1, x2, params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "wehCdvi2ggFp",
    "outputId": "e88315f3-14f5-45f0-b01d-dd680b41c77b"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "1 loop, best of 5: 486 ms per loop\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "%%timeit\n",
    "# Slower - forward pass (FP) is expensive relative to parameters.\n",
    "# Time cost scales poorly with batch size N.\n",
    "ntk_fn_ntvp(x1, x2, params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Yrm53akVggFp",
    "outputId": "13f96e20-25d6-4ecf-d33e-7736e3e02e2b"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "1 loop, best of 5: 185 ms per loop\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "%%timeit\n",
    "# 2X faster!\n",
    "ntk_fn_str_derivatives(x1, x2, params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "P1QtkBqLggFp",
    "outputId": "b7133747-aea9-494b-d293-5a42b4c265b3"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "10 loops, best of 5: 325 ms per loop\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "%%timeit \n",
    "# On TPU should match the fastest method.\n",
    "# On GPU/CPU, currently is broken, and may not be the fastest.\n",
    "ntk_fn_auto(x1, x2, params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t1DLRESpXg7L"
   },
   "source": [
    "# $\\color{blue}O = 128$ logits, batch size $\\color{red}N = 1$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dfoHOoT-gGy6"
   },
   "source": [
    "Both NTK-vector products and Structured derivatives compute NTK faster than Jacobian contraction. NTK-vector products incur no penalty when batch size $\\color{red}N = 1$, and leverage their beneficial scaling with large $\\color{blue}O = 128$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3oHBaAmDhBON",
    "outputId": "3699340f-45c5-49cc-e560-b0e127030028"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "/usr/local/lib/python3.7/dist-packages/neural_tangents/experimental/empirical_tf/empirical.py:232: UserWarning: This function is an early proof-of-concept.\n",
      "  warnings.warn('This function is an early proof-of-concept.')\n"
     ]
    }
   ],
   "source": [
    "O = 128\n",
    "N = 1\n",
    "\n",
    "# Input images x\n",
    "x1 = tf.random.normal((N, *input_shape))\n",
    "x2 = tf.random.normal((N, *input_shape))\n",
    "\n",
    "params, (ntk_fn_jacobian_contraction, ntk_fn_ntvp, ntk_fn_str_derivatives, ntk_fn_auto) = get_ntk_fns(O=O)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "sc1bUvL-KrK9",
    "outputId": "c5b2c294-f7a9-4944-dbe0-d80fd2d3c8c8"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(1, 1, 128, 128)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Jacobian contraction\n",
    "k_1 = ntk_fn_jacobian_contraction(x1, x2, params)\n",
    "print(k_1.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "qdNsmnjOKyp0",
    "outputId": "5e9b611b-d755-4866-ffe4-3e8cc053ad16"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(1, 1, 128, 128)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# NTK-vector products\n",
    "k_2 = ntk_fn_ntvp(x1, x2, params)\n",
    "print(k_2.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Iw6HL260K26E",
    "outputId": "72480346-636f-47f8-b42e-570c96653970"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(1, 1, 128, 128)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Structured derivatives\n",
    "k_3 = ntk_fn_str_derivatives(x1, x2, params)\n",
    "print(k_3.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "DYG0fV9nOjnd",
    "outputId": "3afa8398-2245-4377-d50a-91b733dc3960"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "tf.Tensor(0.016565105, shape=(), dtype=float32) tf.Tensor(0.0032940311, shape=(), dtype=float32) tf.Tensor(0.01386257, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Make sure kernels agree.\n",
    "print(\n",
    "    tf.reduce_max(tf.abs(k_1 - k_2)) / tf.reduce_mean(tf.abs(k_1)), \n",
    "    tf.reduce_max(tf.abs(k_1 - k_3)) / tf.reduce_mean(tf.abs(k_1)),\n",
    "    tf.reduce_max(tf.abs(k_2 - k_3)) / tf.reduce_mean(tf.abs(k_2))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "DyF4M5_HK5Fk",
    "outputId": "85fae406-e1ae-4aa4-a75e-c5bbd4bcf641"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "impl=1, flops=30326192128.0\n",
      "impl=2, flops=25741133824.0\n",
      "impl=3, flops=30259060736.0\n",
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "WARNING:tensorflow:5 out of the last 84 calls to <function convert.<locals>.converted_fun at 0x7f9de1879dd0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "WARNING:tensorflow:5 out of the last 84 calls to <function convert.<locals>.converted_fun at 0x7f9de1879dd0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(1, 1, 128, 128)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Selects best method based on FLOPs at first call / compilation.\n",
    "# Takes about 3x more time to compile.\n",
    "# WARNING: due to an XLA issue, currently only works correctly on TPUs!\n",
    "# Wrong FLOPs for CPU/GPU of JITted functions.\n",
    "k_0 = ntk_fn_auto(x1, x2, params)\n",
    "print(k_0.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "8g1IO71LLJlG",
    "outputId": "87a52a12-7a5a-4334-d059-33ef2db1b6af"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "1 loop, best of 5: 481 ms per loop\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "%%timeit\n",
    "ntk_fn_jacobian_contraction(x1, x2, params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "KEIPcXYRMys2",
    "outputId": "b373236e-3ca3-4048-90cf-c74ff31d2c06"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "1 loop, best of 5: 265 ms per loop\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "%%timeit\n",
    "# 2X faster!\n",
    "ntk_fn_ntvp(x1, x2, params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "gVyUPA8xM1ot",
    "outputId": "46982355-41ca-4436-e5b7-45a516fe0e91"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "1 loop, best of 5: 212 ms per loop\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "%%timeit\n",
    "# 2.5X faster!\n",
    "ntk_fn_str_derivatives(x1, x2, params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "o81r3UHJM34_",
    "outputId": "c0a70d71-5436-4558-f475-423efdcad63c"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "10 loops, best of 5: 264 ms per loop\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "%%timeit \n",
    "# On TPU should match the fastest method.\n",
    "# On GPU/CPU, currently is broken, and may not be the fastest.\n",
    "ntk_fn_auto(x1, x2, params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Dh_aFUZFXm_C"
   },
   "source": [
    "# $\\color{blue}O = 1000$ logits, batch size $\\color{red}N = 1$, full NTK"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8rs4KF07fjjv"
   },
   "source": [
    "Structured derivatives allows to compute full $1000\\times 1000$ ImageNet NTK. Other methods run out of memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Z4MpdSJ7hFsD",
    "outputId": "2e82222e-7c0f-4b41-afbd-30c015882dc5"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "/usr/local/lib/python3.7/dist-packages/neural_tangents/experimental/empirical_tf/empirical.py:232: UserWarning: This function is an early proof-of-concept.\n",
      "  warnings.warn('This function is an early proof-of-concept.')\n"
     ]
    }
   ],
   "source": [
    "O = 1000\n",
    "N = 1\n",
    "\n",
    "# Input images x\n",
    "x1 = tf.random.normal((N, *input_shape))\n",
    "x2 = tf.random.normal((N, *input_shape))\n",
    "\n",
    "params, (ntk_fn_jacobian_contraction, ntk_fn_ntvp, ntk_fn_str_derivatives, ntk_fn_auto) = get_ntk_fns(O=O)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ynBEtkA7d_2G",
    "outputId": "5dde650b-ea83-43af-9465-bd01b6d003a8"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(1, 1, 1000, 1000)\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Structured derivatives - fits in memory!\n",
    "k_3 = ntk_fn_str_derivatives(x1, x2, params)\n",
    "print(k_3.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Afniv9k2d_2H",
    "outputId": "4fbd3d99-7fff-4ced-fd14-59010d83fdec"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "1 loop, best of 5: 1.28 s per loop\n"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "%%timeit\n",
    "ntk_fn_str_derivatives(x1, x2, params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 537
    },
    "id": "4Ou5a6y4zH5Q",
    "outputId": "43cde990-20ad-4f80-94da-fa985ea66cf2"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "error",
     "ename": "UnknownError",
     "evalue": "ignored",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mUnknownError\u001B[0m                              Traceback (most recent call last)",
      "\u001B[0;32m<ipython-input-30-98795d960ac0>\u001B[0m in \u001B[0;36m<module>\u001B[0;34m()\u001B[0m\n\u001B[1;32m      1\u001B[0m \u001B[0;31m# test {\"skip\": true}\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m      2\u001B[0m \u001B[0;31m# NTK-vector products - OOM!\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m----> 3\u001B[0;31m \u001B[0mk_2\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mntk_fn_ntvp\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mx1\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mx2\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mparams\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m      4\u001B[0m \u001B[0mprint\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mk_3\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mshape\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/traceback_utils.py\u001B[0m in \u001B[0;36merror_handler\u001B[0;34m(*args, **kwargs)\u001B[0m\n\u001B[1;32m    151\u001B[0m     \u001B[0;32mexcept\u001B[0m \u001B[0mException\u001B[0m \u001B[0;32mas\u001B[0m \u001B[0me\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    152\u001B[0m       \u001B[0mfiltered_tb\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0m_process_traceback_frames\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0me\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m__traceback__\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 153\u001B[0;31m       \u001B[0;32mraise\u001B[0m \u001B[0me\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mwith_traceback\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mfiltered_tb\u001B[0m\u001B[0;34m)\u001B[0m \u001B[0;32mfrom\u001B[0m \u001B[0;32mNone\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    154\u001B[0m     \u001B[0;32mfinally\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    155\u001B[0m       \u001B[0;32mdel\u001B[0m \u001B[0mfiltered_tb\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/execute.py\u001B[0m in \u001B[0;36mquick_execute\u001B[0;34m(op_name, num_outputs, inputs, attrs, ctx, name)\u001B[0m\n\u001B[1;32m     53\u001B[0m     \u001B[0mctx\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mensure_initialized\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m     54\u001B[0m     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,\n\u001B[0;32m---> 55\u001B[0;31m                                         inputs, attrs, num_outputs)\n\u001B[0m\u001B[1;32m     56\u001B[0m   \u001B[0;32mexcept\u001B[0m \u001B[0mcore\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m_NotOkStatusException\u001B[0m \u001B[0;32mas\u001B[0m \u001B[0me\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m     57\u001B[0m     \u001B[0;32mif\u001B[0m \u001B[0mname\u001B[0m \u001B[0;32mis\u001B[0m \u001B[0;32mnot\u001B[0m \u001B[0;32mNone\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;31mUnknownError\u001B[0m: Failed to determine best cudnn convolution algorithm for:\n%cudnn-conv-bw-filter.15 = (f32[3,3,512,512000]{1,0,2,3}, u8[0]{0}) custom-call(f32[1,7,7,512]{2,1,3,0} %bitcast.1110, f32[1,7,7,512000]{2,1,3,0} %bitcast.2772), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f, custom_call_target=\"__cudnn$convBackwardFilter\", metadata={op_type=\"XlaConvV2\" op_name=\"XlaConvV2_109\" source_file=\"/usr/local/lib/python3.7/dist-packages/tensorflow/compiler/tf2xla/ops/gen_xla_ops.py\" source_line=478}, backend_config=\"{\\\"conv_result_scale\\\":1,\\\"activation_mode\\\":\\\"0\\\",\\\"side_input_scale\\\":0}\"\n\nOriginal error: RESOURCE_EXHAUSTED: Allocating 9437184000 bytes exceeds the memory limit of 4294967296 bytes.\n\nTo ignore this failure and try to use a fallback algorithm (which may have suboptimal performance), use XLA_FLAGS=--xla_gpu_strict_conv_algorithm_picker=false.  Please also file a bug for the root cause of failing autotuning. [Op:__inference_converted_fun_1957800]"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# NTK-vector products - OOM!\n",
    "k_2 = ntk_fn_ntvp(x1, x2, params)\n",
    "print(k_3.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 537
    },
    "id": "rjXRalQfd_2G",
    "outputId": "fe6e5f37-1b3c-4665-e24b-79c2700c0cae"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
     ]
    },
    {
     "output_type": "error",
     "ename": "UnknownError",
     "evalue": "ignored",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mUnknownError\u001B[0m                              Traceback (most recent call last)",
      "\u001B[0;32m<ipython-input-31-628f1776b6c6>\u001B[0m in \u001B[0;36m<module>\u001B[0;34m()\u001B[0m\n\u001B[1;32m      1\u001B[0m \u001B[0;31m# test {\"skip\": true}\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m      2\u001B[0m \u001B[0;31m# Jacobian contraction - OOM!\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m----> 3\u001B[0;31m \u001B[0mk_1\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mntk_fn_jacobian_contraction\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mx1\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mx2\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mparams\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m      4\u001B[0m \u001B[0mprint\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mk_1\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mshape\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/traceback_utils.py\u001B[0m in \u001B[0;36merror_handler\u001B[0;34m(*args, **kwargs)\u001B[0m\n\u001B[1;32m    151\u001B[0m     \u001B[0;32mexcept\u001B[0m \u001B[0mException\u001B[0m \u001B[0;32mas\u001B[0m \u001B[0me\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    152\u001B[0m       \u001B[0mfiltered_tb\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0m_process_traceback_frames\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0me\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m__traceback__\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 153\u001B[0;31m       \u001B[0;32mraise\u001B[0m \u001B[0me\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mwith_traceback\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mfiltered_tb\u001B[0m\u001B[0;34m)\u001B[0m \u001B[0;32mfrom\u001B[0m \u001B[0;32mNone\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    154\u001B[0m     \u001B[0;32mfinally\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    155\u001B[0m       \u001B[0;32mdel\u001B[0m \u001B[0mfiltered_tb\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/execute.py\u001B[0m in \u001B[0;36mquick_execute\u001B[0;34m(op_name, num_outputs, inputs, attrs, ctx, name)\u001B[0m\n\u001B[1;32m     53\u001B[0m     \u001B[0mctx\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mensure_initialized\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m     54\u001B[0m     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,\n\u001B[0;32m---> 55\u001B[0;31m                                         inputs, attrs, num_outputs)\n\u001B[0m\u001B[1;32m     56\u001B[0m   \u001B[0;32mexcept\u001B[0m \u001B[0mcore\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m_NotOkStatusException\u001B[0m \u001B[0;32mas\u001B[0m \u001B[0me\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m     57\u001B[0m     \u001B[0;32mif\u001B[0m \u001B[0mname\u001B[0m \u001B[0;32mis\u001B[0m \u001B[0;32mnot\u001B[0m \u001B[0;32mNone\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;31mUnknownError\u001B[0m: Failed to determine best cudnn convolution algorithm for:\n%cudnn-conv-bw-filter.1 = (f32[3,3,512,512000]{1,0,2,3}, u8[0]{0}) custom-call(f32[1,7,7,512]{2,1,3,0} %bitcast.2270, f32[1,7,7,512000]{2,1,3,0} %bitcast.5528), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f, custom_call_target=\"__cudnn$convBackwardFilter\", metadata={op_type=\"XlaConvV2\" op_name=\"XlaConvV2_213\" source_file=\"/usr/local/lib/python3.7/dist-packages/tensorflow/compiler/tf2xla/ops/gen_xla_ops.py\" source_line=478}, backend_config=\"{\\\"conv_result_scale\\\":1,\\\"activation_mode\\\":\\\"0\\\",\\\"side_input_scale\\\":0}\"\n\nOriginal error: RESOURCE_EXHAUSTED: Allocating 9437184000 bytes exceeds the memory limit of 4294967296 bytes.\n\nTo ignore this failure and try to use a fallback algorithm (which may have suboptimal performance), use XLA_FLAGS=--xla_gpu_strict_conv_algorithm_picker=false.  Please also file a bug for the root cause of failing autotuning. [Op:__inference_converted_fun_2087253]"
     ]
    }
   ],
   "source": [
    "# test {\"skip\": true}\n",
    "# Jacobian contraction - OOM!\n",
    "k_1 = ntk_fn_jacobian_contraction(x1, x2, params)\n",
    "print(k_1.shape)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "empirical_ntk_resnet_tf.ipynb",
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
